import requests
import json
from PIL import Image
import io
import base64
import os
import concurrent.futures
import time
import threading

# 创建线程安全的打印锁
print_lock = threading.Lock()

def print_png_files(directory, max_workers=8):
    """
    使用多线程处理目录下所有.jpg文件
    
    参数:
        directory (str): 要搜索的目录路径
        max_workers (int): 最大线程数
    """
    jpg_files = []
    
    # 收集所有.jpg文件路径
    with print_lock:
        print(f"扫描目录: {directory}")
    
    for root, _, files in os.walk(directory):
        for file in files:
            if file.lower().endswith('.jpg'):
                full_path = os.path.join(root, file)
                jpg_files.append(full_path)
    
    with print_lock:
        print(f"找到 {len(jpg_files)} 个JPG文件")
        print("开始并行处理...")
    
    # 使用线程池处理文件
    start_time = time.time()
    
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        # 提交所有任务到线程池
        future_to_file = {executor.submit(call_detect_api, file): file for file in jpg_files}
        
        # 处理完成的任务
        for future in concurrent.futures.as_completed(future_to_file):
            file = future_to_file[future]
            try:
                result = future.result()
                if result is not None:
                    with print_lock:
                        print(f"✓ 成功处理: {file}")
                else:
                    with print_lock:
                        print(f"✗ 处理失败: {file}")
            except Exception as e:
                with print_lock:
                    print(f"⚠ 处理异常: {file}, 错误: {str(e)}")
    
    elapsed = time.time() - start_time
    with print_lock:
        print(f"\n处理完成! 共处理 {len(jpg_files)} 个文件")
        print(f"总耗时: {elapsed:.2f}秒 | 平均: {len(jpg_files)/max(1, elapsed):.2f}文件/秒")

def call_detect_api(image_path):
    """
    调用AI检测API并将返回的JSON结果保存到文件
    
    参数:
        image_path (str): 图片文件的完整路径
    
    返回:
        dict: API返回的JSON数据，如果出错返回None
    """
    try:
        # 打开图片并转换为JPEG格式的Base64编码
        with Image.open(image_path) as img:
            buffered = io.BytesIO()
            img.convert('RGB').save(buffered, format='JPEG')
            base64_encoded = base64.b64encode(buffered.getvalue()).decode('utf-8')
        
        # 构造请求负载
        payload = {"image": {"file": base64_encoded}}
        url = "http://192.168.2.230:3154/x-api/v1/ai/detect/all"
        headers = {"Content-Type": "application/json"}
        
        # 发送POST请求
        response = requests.post(
            url, 
            headers=headers, 
            data=json.dumps(payload),
            timeout=15  # 增加超时时间
        )
        
        # 检查HTTP状态码
        response.raise_for_status()
        
        # 解析JSON响应
        result_json = response.json()
        
        # 生成输出文件路径
        directory = os.path.dirname(image_path)
        filename = os.path.splitext(os.path.basename(image_path))[0] + '.json'
        output_path = os.path.join(directory, filename)
        
        # 将JSON结果写入文件
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(result_json, f, indent=2, ensure_ascii=False)
        
        return result_json
        
    except Exception as e:
        with print_lock:
            print(f"处理图片时出错: {image_path}")
            print(f"错误详情: {str(e)}")
        return None

if __name__ == "__main__":
    # 替换为你的实际图片路径
    image_dir = "./pic"
    
    # 设置线程数 (根据网络和CPU调整)
    max_workers = 38
    
    print(f"启动多线程处理 (线程数: {max_workers})")
    print("=" * 50)
    print_png_files(image_dir, max_workers=max_workers)

